import sys
import os
import ast
import io
import re
import tokenize
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, List, Optional, Set, Tuple, Union
from PyQt5.QtWidgets import (QApplication, QMainWindow, QVBoxLayout,
                             QHBoxLayout, QWidget, QPushButton, QFileDialog,
                             QLineEdit, QTextEdit, QLabel, QProgressBar,
                             QCheckBox, QGroupBox, QMessageBox)
from PyQt5.QtCore import QThread, pyqtSignal
from PyQt5.QtGui import QFont


@dataclass(frozen=True)
class _Span:
    """一个源码区间，以 (start_line, start_col) 和 (end_line, end_col) 定位"""
    start: Tuple[int, int]
    end: Tuple[int, int]


class CommentRemover:
    """Python代码注释清理器"""

    def __init__(self):
        self.preserve_docstrings = True
        self.preserve_shebang = True
        self.preserve_encoding = True

        # PEP 263 编码声明
        self._coding_re = re.compile(r"^[ \t\f]*#.*?coding[:=][ \t]*([-\w.]+)")

    # --------------- 公共 API ---------------
    def clean_code(self, code: str) -> str:
        """传入源码字符串，返回移除注释后的源码字符串。"""
        docstring_spans = self._collect_docstring_spans(code) if not self.preserve_docstrings else set()
        tokens = tokenize.tokenize(io.BytesIO(code.encode("utf-8")).readline)

        out_tokens: List[tokenize.TokenInfo] = []
        for tok in tokens:
            ttype = tok.type
            tstr = tok.string
            # 跳过 ENCODING 伪token（untokenize 不需要它）
            if ttype == tokenize.ENCODING:
                continue

            # 处理注释（包括 shebang 与编码声明）
            if ttype == tokenize.COMMENT:
                if self._should_keep_comment(tok):
                    out_tokens.append(tok)
                # 否则丢弃该注释
                continue

            # 处理 docstring：仅当其是 AST 判定的 docstring 时才删除
            if not self.preserve_docstrings and ttype == tokenize.STRING:
                if self._is_docstring_token(tok, docstring_spans):
                    # 注意：如果该 docstring 是块内唯一语句，删除后代码块会为空；
                    # 如需保持语法有效，可在需要时手动插入 'pass'（此实现按需求不强制插入）。
                    continue

            out_tokens.append(tok)

        # 回拼源码。untokenize 会保持原始布局（行内空白、缩进等）
        cleaned = tokenize.untokenize(out_tokens)
        # 注：删除行尾注释后，untokenize 可能留下对应位置的空白，以保证列位置一致；
        # 这不会影响缩进与语义。如需去掉行尾多余空白，可自行对每行 rstrip。
        return cleaned

    @staticmethod
    def clean_consecutive_empty_lines(text):
        """
        删除连续的空行，只保留一个空行
        空行定义：空字符串或只包含空格/tab字符的行
        """
        # 匹配连续的空行（包括只有空格/tab的行）
        # \n[ \t]*\n 匹配一个空行
        # ([ \t]*\n)* 匹配后续连续的空行
        return re.sub(r'\n[ \t]*\n([ \t]*\n)*', '\n\n', text)

    @staticmethod
    def remove_consecutive_empty_lines_clear(text):
        lines = text.split('\n')
        result = []
        empty_line_count = 0

        for line in lines:
            # 检查是否为空行（空字符串或只包含空格/tab）
            if line.strip() == '':
                empty_line_count += 1
                if empty_line_count == 1:
                    result.append('')  # 添加一个纯空行
            else:
                # 非空行
                empty_line_count = 0
                result.append(line)

        return '\n'.join(result)

    def clean_file(
        self,
        input_path: Union[str, Path],
        output_path: Optional[Union[str, Path]] = None,
        inplace: bool = False
    ) -> str:
        """
        读取文件，返回清理后的代码。
        - inplace=True 时会覆盖原文件；
        - 或指定 output_path 写出到新文件。
        """
        # tokenize.open 会根据文件内编码声明解码
        with tokenize.open(str(input_path)) as f:
            code = f.read()

        cleaned = self.clean_code(code)
        cleaned = CommentRemover.clean_consecutive_empty_lines(cleaned)

        if inplace:
            Path(input_path).write_text(cleaned, encoding="utf-8")
        elif output_path is not None:
            Path(output_path).write_text(cleaned, encoding="utf-8")

        return cleaned

    # --------------- 内部工具 ---------------
    def _should_keep_comment(self, tok: tokenize.TokenInfo) -> bool:
        """根据配置判断该 COMMENT token 是否保留（shebang/编码声明）。"""
        s = tok.string
        line_no, col = tok.start

        # shebang 只在首行且首列
        if line_no == 1 and col == 0 and s.startswith("#!"):
            return self.preserve_shebang

        # 编码声明只在前两行有效（PEP 263）
        if line_no <= 2 and self._coding_re.match(s):
            return self.preserve_encoding

        # 其他注释均删除
        return False

    def _collect_docstring_spans(self, code: str) -> Set[_Span]:
        """用 AST 精确收集所有 docstring 的源码区间。"""
        spans: Set[_Span] = set()
        try:
            tree = ast.parse(code)
        except SyntaxError:
            # 无法解析时，不做 docstring 删除（以免误删普通字符串）
            return spans

        def maybe_add(node: ast.AST):
            body = getattr(node, "body", None)
            if not body:
                return
            first = body[0]
            # 兼容 3.8+（ast.Constant），老版本可扩展添加 ast.Str 判断
            if isinstance(first, ast.Expr) and isinstance(getattr(first, "value", None), ast.Constant) and isinstance(first.value.value, str):
                # 需要 Python 3.8+ 的 end_lineno/end_col_offset 才能精确定位
                if hasattr(first, "end_lineno") and hasattr(first, "end_col_offset"):
                    spans.add(_Span(
                        start=(first.lineno, first.col_offset),
                        end=(first.end_lineno, first.end_col_offset),
                    ))

        # 模块级
        maybe_add(tree)
        # 函数/类级
        for n in ast.walk(tree):
            if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                maybe_add(n)

        return spans

    @staticmethod
    def _is_docstring_token(tok: tokenize.TokenInfo, spans: Set[_Span]) -> bool:
        """判断一个 STRING token 是否恰好是 AST 标注的 docstring。"""
        if not spans:
            return False
        start = tok.start
        end = tok.end
        for sp in spans:
            if sp.start == start and sp.end == end:
                return True
        return False


class ProcessThread(QThread):
    """后台处理线程"""
    progress_updated = pyqtSignal(int, str)
    file_processed = pyqtSignal(str, bool, str)
    finished_processing = pyqtSignal(list)

    def __init__(self, paths, is_directory=False):
        super().__init__()
        self.paths = paths
        self.is_directory = is_directory
        self.comment_remover = CommentRemover()

    def run(self):
        results = []

        if self.is_directory:
            # 处理目录中的所有Python文件
            all_files = []
            for path in self.paths:
                for root, dirs, files in os.walk(path):
                    for file in files:
                        if file.endswith('.py'):
                            all_files.append(os.path.join(root, file))
        else:
            # 处理选定的文件
            all_files = self.paths

        total_files = len(all_files)

        for i, file_path in enumerate(all_files):
            self.progress_updated.emit(int((i / total_files) * 100),
                                       f"处理中: {os.path.basename(file_path)}")

            output_path = file_path.replace(".py", "_release.py")
            res_str = self.comment_remover.clean_file(file_path, output_path)
            error = ''
            success = True

            result = {
                'input_file': file_path,
                'output_file': output_path,
                'success': success,
                'error': error
            }
            results.append(result)

            self.file_processed.emit(file_path, success, error)

        self.progress_updated.emit(100, "处理完成")
        self.finished_processing.emit(results)


class PythonCleanerGUI(QMainWindow):
    """主窗口"""

    def __init__(self):
        super().__init__()
        self.initUI()
        self.process_thread = None

    def initUI(self):
        self.setWindowTitle('Python代码清理工具')
        self.setGeometry(100, 100, 800, 600)

        # 创建中心部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        # 主布局
        main_layout = QVBoxLayout(central_widget)

        # 文件选择区域
        file_group = QGroupBox("选择文件或目录")
        file_layout = QVBoxLayout(file_group)

        # 目录选择
        dir_layout = QHBoxLayout()
        self.dir_path = QLineEdit()
        self.dir_path.setPlaceholderText("选择项目目录...")
        dir_btn = QPushButton("选择目录")
        dir_btn.clicked.connect(self.select_directory)
        dir_layout.addWidget(self.dir_path)
        dir_layout.addWidget(dir_btn)

        # 文件选择
        file_layout_h = QHBoxLayout()
        self.file_paths = QLineEdit()
        self.file_paths.setPlaceholderText("或选择具体的Python文件...")
        file_btn = QPushButton("选择文件")
        file_btn.clicked.connect(self.select_files)
        file_layout_h.addWidget(self.file_paths)
        file_layout_h.addWidget(file_btn)

        file_layout.addLayout(dir_layout)
        file_layout.addLayout(file_layout_h)

        # 选项设置
        options_group = QGroupBox("处理选项")
        options_layout = QVBoxLayout(options_group)

        self.preserve_docstrings = QCheckBox("保留文档字符串")
        self.preserve_docstrings.setChecked(True)
        self.preserve_shebang = QCheckBox("保留shebang行 (#!/usr/bin/env python)")
        self.preserve_shebang.setChecked(True)
        self.preserve_encoding = QCheckBox("保留编码声明")
        self.preserve_encoding.setChecked(True)

        options_layout.addWidget(self.preserve_docstrings)
        options_layout.addWidget(self.preserve_shebang)
        options_layout.addWidget(self.preserve_encoding)

        # 控制按钮
        control_layout = QHBoxLayout()
        self.start_btn = QPushButton("开始处理")
        self.start_btn.clicked.connect(self.start_processing)
        self.start_btn.setStyleSheet("QPushButton { background-color: #4CAF50; color: white; font-weight: bold; }")
        control_layout.addWidget(self.start_btn)
        control_layout.addStretch()

        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_label = QLabel("就绪")

        # 结果显示区域
        self.result_text = QTextEdit()
        self.result_text.setFont(QFont("Consolas", 9))
        self.result_text.setReadOnly(True)

        # 添加到主布局
        main_layout.addWidget(file_group)
        main_layout.addWidget(options_group)
        main_layout.addLayout(control_layout)
        main_layout.addWidget(self.progress_label)
        main_layout.addWidget(self.progress_bar)
        main_layout.addWidget(QLabel("处理结果:"))
        main_layout.addWidget(self.result_text)

    def select_directory(self):
        directory = QFileDialog.getExistingDirectory(self, "选择项目目录")
        if directory:
            self.dir_path.setText(directory)
            self.file_paths.clear()

    def select_files(self):
        files, _ = QFileDialog.getOpenFileNames(
            self, "选择Python文件", "", "Python Files (*.py)")
        if files:
            self.file_paths.setText("; ".join(files))
            self.dir_path.clear()

    def start_processing(self):
        # 获取输入路径
        directory = self.dir_path.text().strip()
        files = self.file_paths.text().strip()

        if not directory and not files:
            QMessageBox.warning(self, "警告", "请选择目录或文件")
            return

        # 禁用开始按钮
        self.start_btn.setEnabled(False)
        self.result_text.clear()

        # 设置处理选项
        if hasattr(self, 'process_thread') and self.process_thread:
            self.process_thread.quit()
            self.process_thread.wait()

        # 创建处理线程
        if directory:
            self.process_thread = ProcessThread([directory], is_directory=True)
        else:
            file_list = [f.strip() for f in files.split(";") if f.strip()]
            self.process_thread = ProcessThread(file_list, is_directory=False)

        # 连接信号
        self.process_thread.progress_updated.connect(self.update_progress)
        self.process_thread.file_processed.connect(self.file_processed)
        self.process_thread.finished_processing.connect(self.processing_finished)

        # 启动线程
        self.process_thread.start()

    def update_progress(self, value, text):
        self.progress_bar.setValue(value)
        self.progress_label.setText(text)

    def file_processed(self, input_file, success, error):
        if success:
            self.result_text.append(f"✓ 成功处理: {input_file}")
        else:
            self.result_text.append(f"✗ 处理失败: {input_file}")
            self.result_text.append(f"  错误: {error}")

    def processing_finished(self, results):
        success_count = sum(1 for r in results if r['success'])
        total_count = len(results)

        self.result_text.append("\n" + "=" * 50)
        self.result_text.append(f"处理完成: {success_count}/{total_count} 个文件成功处理")

        if success_count > 0:
            self.result_text.append("\n生成的文件:")
            for result in results:
                if result['success']:
                    self.result_text.append(f"  {result['output_file']}")

        self.start_btn.setEnabled(True)
        self.progress_label.setText("就绪")


def main():
    app = QApplication(sys.argv)
    window = PythonCleanerGUI()
    window.show()
    sys.exit(app.exec_())


if __name__ == '__main__':
    main()